# -*- coding: utf-8 -*-
# @Time : 2021/11/1 15:37
# @Author : Ming
# @FileName: embedding_linear_layer.py
# @Software: PyCharm

import torch

'''
由于one-hot编码过于稀疏，高维度，硬编码，所以加以进化，使用嵌入层（embedding layer）

'''

# parameters
num_class = 4
input_size = 4
hidden_size = 8
embedding_size = 10
num_layers = 2
batch_size = 1
seq_len = 5

# 通过RNN实现hello转换为ohlol
idx2char = ['e', 'h', 'l', 'o']
x_data = [[1, 0, 2, 2, 3]] # (batch, seq_len),这里变成了2个维度
y_data = [3, 1, 2, 3, 2] # (batch * seq_len)

inputs = torch.LongTensor(x_data)   # (batch, seq_len),这里变成了2个维度
labels = torch.LongTensor(y_data)   # (batch * seq_len)

# 设计模型
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.emb = torch.nn.Embedding(input_size, embedding_size)   #嵌入层，软编码，稠密
        self.rnn = torch.nn.RNN(input_size=embedding_size, # 输入维度和嵌入层维度一致
                                hidden_size=hidden_size,    # 隐层
                                num_layers=num_layers,  #RNN层数
                                batch_first=True) # 是否是批量（batch）在前
        self.fc = torch.nn.Linear(hidden_size, num_class)

    def forward(self, x):
        hidden = torch.zeros(num_layers, x.size(0), hidden_size)
        x = self.emb(x) # (batch, seqLen, embeddingSize)
        x, _ = self.rnn(x, hidden)
        x = self.fc(x)
        return x.view(-1, num_class)


net = Model()
criterion = torch.nn.CrossEntropyLoss() # 交叉熵
optimizer = torch.optim.Adam(net.parameters(), lr=0.05) # 优化器

# 训练循环开始
for epoch in range(15):
    optimizer.zero_grad()
    # RNN是直接获取所有输入数据
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('Predicted: ', ''.join([idx2char[x] for x in idx]), end='')
    print(', Epoch [%d/15] loss = %.3f' % (epoch + 1, loss.item()))